import re
import pandas as pd
def extract_code_info(code: str):
    function_pattern = re.compile(r'\b(def|function|public|private|protected|static)\s+(\w+)\s?\(')
    class_pattern = re.compile(r'\b(class|interface|abstract)\s+(\w+)')
    single_line_comment_pattern = re.compile(r'//.*|#.*')
    multi_line_comment_pattern = re.compile(r'/\*.*?\*/', re.DOTALL)
    python_docstring_pattern = re.compile(r'\"\"\".*?\"\"\"', re.DOTALL)
    

    functions = [match.group(2) for match in function_pattern.finditer(code)]
    classes = [match.group(2) for match in class_pattern.finditer(code)]
    single_line_comments = single_line_comment_pattern.findall(code)
    multi_line_comments = multi_line_comment_pattern.findall(code)
    python_docstrings = python_docstring_pattern.findall(code)
    
    result = {
        'functions': functions,
        'classes': classes,
        'comments': single_line_comments + multi_line_comments + python_docstrings
    }
    
    return result

def extract_code_lines(code: str):
    function_pattern = re.compile(r'\b(def|function|public|private|protected|static)\s+(\w+)\s?\(')
    
    class_pattern = re.compile(r'\b(class|interface|abstract)\s+(\w+)')
    
    single_line_comment_pattern = re.compile(r'//.*|#.*')
    multi_line_comment_pattern = re.compile(r'/\*.*?\*/', re.DOTALL)
    
    functions = []
    classes = []
    comments = []
    
    lines = code.splitlines()
    
    for line in lines:
        if function_pattern.search(line):
            functions.append(line.strip())
        if class_pattern.search(line):
            classes.append(line.strip())
        if single_line_comment_pattern.search(line) or multi_line_comment_pattern.search(line):
            comments.append(line.strip())
    
    result = {
        'functions': functions,
        'classes': classes,
        'comments': comments
    }
    
    return result
def load_data(input_filepath = "Q_B_without_answer.jsonl"):
    df = pd.read_json(input_filepath,lines=True)
    prefixs = df["prefix"].tolist()
    suffixs = df["fim_suffix"].tolist()
    return prefixs, suffixs
def get_first_line(text: str) -> str:
    lines = text.splitlines()
    for line in lines:
        if line.strip():  
            return line
    return ""  

def get_last_line(text: str) -> str:
    lines = text.splitlines()
    for line in reversed(lines):
        if line.strip():  
            return line
    return ""  




def extract_middle_string(prefix, suffix, string):
    escaped_prefix = re.escape(prefix)
    escaped_suffix = re.escape(suffix)
    if suffix == "":
        pattern = re.compile(r'{}(.*)'.format(escaped_prefix), re.DOTALL)
    else:
        pattern = re.compile(r'{}(.*?)(?={})'.format(escaped_prefix, escaped_suffix), re.DOTALL)
    match = pattern.search(string)
    if match:
        return match.group(1)
    else:
        return None
      
if __name__ == "__main__":
    pass
    